# for handle netCDF4 data
import netCDF4 as nc
import numpy as np
import dask.array as da
from prettytable import PrettyTable
import re
import pandas as pd

def _str_some_value(array: np.ndarray):
    # if less than 10, print all, else print first 5 and last 5
    if np.shape(array)[0] <= 10:
        return str(array)
    else:
        return str(array)
    
def _convert_time_unit(unit:str):
    """Convert time unit to numpy time unit"""
    np_unit = {
        "Y": ["years", "year", "Y"],
        "M": ["months", "month", "M"],
        "D": ["days", "day", "D"],
        "h": ["hours", "hour", "h"],
        "m": ["minutes", "minute", "m"],
        "s": ["seconds", "second", "s"],
        "ms": ["milliseconds", "millisecond", "ms"],
    }
    for key, value in np_unit.items():
        if unit in value:
            return key
    raise ValueError(f"Can't convert unit {unit} to numpy time unit")

# dim
class Dimension(object):
    """For 1-D coordinate variables"""
    isTimeDimension = False
    
    # if isTimeDimension, datetime is not None, and use datetime as values
    datetime = None

    def set_nc_dim(self, dim:nc._netCDF4.Dimension):
        self.dim = dim
        self.values = np.array(dim[:])
        # check if is time dimension by units like "xxx since xxx"
        if hasattr(dim, "units"):
            units = dim.units
            if units.find("since") != -1:
                self.isTimeDimension = True
                self._convert_time_dimension()

    def _convert_time_dimension(self):
        # units like "days since 1979-01-01 00:00:00"
        # convert to days since 1979-01-01 00:00:00
        units = self.dim.units
        unit, startTime = re.findall(r"(.*?)\ssince\s(.*)", units)[0]
        
        startTime = np.datetime64(startTime)
        np_unit = _convert_time_unit(unit)

        self.values = startTime + np.array(self.dim[:], dtype=f"timedelta64[{np_unit}]")


    @property
    def name(self):
        return self.dim.name
    
    @property
    def shape(self):
        return self.values.shape
    
    @property
    def units(self):
        return self.dim.units

    def preview(self):
        if type(self.values) == np.ndarray:
            threshold = np.get_printoptions()['threshold']
            np.set_printoptions(threshold=10)
            s = np.array_str(self.values)
            np.set_printoptions(threshold=threshold)
            return s
        else:
            return str(self.values)

    def __repr__(self) -> str:
        s = f"<Dimension: {self.name} (shape: {self.shape}, dtype: {self.dim.dtype})>" 
        if self.shape[0] < 1000:
            s += "\n" + self.preview()
        return s

    def getIndex(self, value):
        """Get index of value in values"""
        if type(value) == str and self.isTimeDimension:
            value = np.datetime64(value)
        try:
            _index = np.where(self.values == value)[0][0]
            return _index
        except:
            raise ValueError(f"No such value: {value} in '{self.name}' dimension.") 

    def getListIndex(self, values):
        """Get index of values in values"""
        if type(values) == str and self.isTimeDimension:
            values = np.datetime64(values)
        try:
            _index = np.where(np.isin(self.values, values))[0]
            return _index
        except:
            raise ValueError(f"No such value: {values} in '{self.name}' dimension.")

    def select(self, indices):
        self.values = self.values[indices]

# dataArray

class Variable(object):
    """等于 xarray.DataArray"""
    def set_nc_var(self, var:nc._netCDF4.Variable, dims:tuple[Dimension]):
        self.name = var.name
        self.dtype = var.dtype
        self.dims = dims
        self.dimNames = tuple([dim.name for dim in dims])
        for dim in dims:
            setattr(self, dim.name, dim)
        self.values = da.from_array(var)
    
    @property
    def shape(self):
        return self.values.shape

    def __repr__(self) -> str:
        pt = PrettyTable(["Dimension", "Shape", "dtype", "Preview"])
        pt.align["Dimension"] = "l"
        pt.align["Preview"] = "l"
        if len(self.shape) == 0:
            pt.add_row(['*'+self.name, '-', 
                        self.dtype, self.values])
        else:
            pt.add_row(['*'+self.name, self.shape, self.dtype, self.dimNames])
        for dim in self.dims:
            pt.add_row([dim.name, dim.shape, dim.dim.dtype, dim.preview()])
        return pt.get_string()

    def _get_dimension_index(self, dimName:str):
        return self.dimNames.index(dimName)
    
    def sel(self, **kwargs):
        """Select data by dimension name and value"""

        indices = []
        for i in range(len(self.dimNames)):
            hasArgsFlat = False
            for key, value in kwargs.items():
                if key == self.dimNames[i]:
                    hasArgsFlat = True
                    # 如果是单值
                    if type(value) == int or type(value) == float or type(value) == str:
                        _index = self.dims[i].getIndex(value)
                        indices.append(_index)
                        self.dims[i].select(_index)

                    # 如果是 slice
                    elif type(value) == slice:
                        start = self.dims[i].getIndex(value.start)
                        stop = self.dims[i].getIndex(value.stop)
                        indices.append(slice(start, stop+1, value.step))
                        self.dims[i].select(slice(start, stop+1, value.step))
                    
                    else:
                        _index = self.dims[i].getListIndex(value)
                        indices.append(_index)
                        self.dims[i].select(_index)
            
                    break

            if not hasArgsFlat:
                indices.append(slice(None, None))
        
        self.values = self.values[tuple(indices)]            
        # remove void dimension
        self.dims = tuple([dim for dim in self.dims if len(dim.shape) > 0])
        self.dimNames = tuple([dim.name for dim in self.dims])

        return self
    
    def _remove_dim(self, dimName:str):
        """Remove dimension by name"""
        self.dims = tuple([dim for dim in self.dims if dim.name != dimName])
        self.dimNames = tuple([dim.name for dim in self.dims])
        
    def mean(self, dim:str | list[str] = None):
        """Calculate mean value along dim"""
        if dim is None:
            self.values = self.values.mean()
            # remove all dimension
            self.dims = tuple()
            self.dimNames = tuple()
        else:
            if type(dim) == str:
                dim = [dim]
            dimIndex = [self._get_dimension_index(dimName) for dimName in dim]
            self.values = self.values.mean(axis=dimIndex)
            for dimName in dim:
                self._remove_dim(dimName)
        return self

    def compute(self):
        self.values = self.values.compute()
        return self

class LiData(object):
    def __init__(self) -> None:
        pass

    def _get_dim(self, dimName:tuple[str]):
        res = {}
        for name in dimName:
            res[name] = self.dataset.variables[name]
        return res

    def load(self, path:str):
        self.dataset = nc.Dataset(path)
        dim = self.dataset.dimensions
        var = self.dataset.variables

        varNotInDim = []
        varInDim = []
        for key in var.keys():
            if key not in dim.keys():
                varNotInDim.append(key)
            else:
                varInDim.append(key)

        # 把每一个变量都存储为类的属性，可以直接用.访问，如：data.u
        # 设置 dimension
        for key in varInDim:
            _dim = Dimension()
            _dim.set_nc_dim(var[key])
            setattr(self, key, _dim)
        
        # 设置 variables
        for key in varNotInDim:
            _var = Variable()
            _var.set_nc_var(
                var[key], 
                tuple([getattr(self, dimName) for dimName in var[key].dimensions])
            )
            setattr(self, key, _var)


def load_data(path:str):
    ld = LiData()
    ld.load(path)
    return ld

def sel_nc(varData, axis:int, axisValues:np.ndarray, value:int | float | slice | np.ndarray):
    """Select data by axis and value"""
    if type(value) == int or type(value) == float:
        _index = np.where(axisValues == value)[0][0]
        axisValues = axisValues[_index]
        varData = varData.take(_index, axis=axis)
    elif type(value) == slice:
        start = np.where(axisValues == value.start)[0][0]
        stop = np.where(axisValues == value.stop)[0][0]
        axisValues = axisValues[start:stop+1]
        varData = varData.take(np.arange(start, stop+1), axis=axis)
    else:
        _index = np.where(np.isin(axisValues, value))[0]
        axisValues = axisValues[_index]
        varData = varData.take(_index, axis=axis)
    return varData, axisValues

if __name__ == "__main__":
    d = load_data("/Users/jjli/research/cache/BPF_multilevel/uwnd_BPF_3-8_multilevel_daily.nc")
    # d = load_data("/Datadisk/ERA5/4xdaily0.5/uwnd/uwnd.197902.nc")
    d = d.u.sel(level=850, 
                  longitude=slice(120, 130), 
                  latitude=slice(10, 0),
                  time=slice("2000-02-01", "2000-02-10"))
    mean = d.mean(dim=['longitude', 'latitude']).compute()
    print(mean)